package idset

import (
	"sort"
	"sync"
	"sync/atomic"
	"time"
)

// idNode 带读写锁
type idNode struct {
	sync.RWMutex
	fromSet  map[string]struct{}
	accessAt int64 // unix nano
}

func (n *idNode) touch(now int64) {
	atomic.StoreInt64(&n.accessAt, now)
}

func (n *idNode) expired(now int64, ttl time.Duration) bool {
	return now-atomic.LoadInt64(&n.accessAt) > int64(ttl)
}

// Manager
type Manager struct {
	data     sync.Map  // ID(string) -> *idNode
	listPool sync.Pool // 复用 []string
}

func NewManager() *Manager {
	m := &Manager{}
	m.listPool.New = func() interface{} {
		return make([]string, 0, 1024)
	}
	return m
}

// Range 遍历所有ID及其对应的from列表（去重且有序）
// 回调函数返回false时停止遍历
func (m *Manager) Range(f func(id string, set []string) bool) {
	m.data.Range(func(key, value interface{}) bool {
		id := key.(string)
		node := value.(*idNode)

		// 复用GetFromList的逻辑获取去重有序列表
		list := m.listPool.Get().([]string)[:0]
		node.RLock()
		for f := range node.fromSet {
			list = append(list, f)
		}
		node.RUnlock()
		sort.Strings(list)

		// 传递副本给回调，避免外部修改影响内部
		ret := make([]string, len(list))
		copy(ret, list)
		m.listPool.Put(list)

		// 执行回调，根据返回值决定是否继续遍历
		return f(id, ret)
	})
}

// Add 添加一个 From；nowTime 由调用者传入
func (m *Manager) Add(id string, from string, nowTime time.Time) (added bool) {
	now := nowTime.UnixNano()
	actual, _ := m.data.LoadOrStore(id, &idNode{
		fromSet: make(map[string]struct{}),
	})
	node := actual.(*idNode)
	node.touch(now)
	// 快速判断，不存在再加写锁
	node.RLock()
	_, ok := node.fromSet[from]
	node.RUnlock()
	if !ok {
		node.Lock()
		if _, ok = node.fromSet[from]; !ok {
			node.fromSet[from] = struct{}{}
			added = true
		}
		node.Unlock()
	}
	return
}

// Adds 批量添加
func (m *Manager) Adds(id string, fromList []string, nowTime time.Time) {
	now := nowTime.UnixNano()
	actual, _ := m.data.LoadOrStore(id, &idNode{
		fromSet: make(map[string]struct{}),
	})
	node := actual.(*idNode)
	node.touch(now)
	node.Lock()
	for _, f := range fromList {
		node.fromSet[f] = struct{}{}
	}
	node.Unlock()
}

// GetFromList 返回去重且有序的 From 列表（只读）
func (m *Manager) GetFromList(id string) []string {
	v, ok := m.data.Load(id)
	if !ok {
		return nil
	}
	node := v.(*idNode)

	list := m.listPool.Get().([]string)[:0]
	node.RLock()
	for f := range node.fromSet {
		list = append(list, f)
	}
	node.RUnlock()
	sort.Strings(list)
	ret := make([]string, len(list))
	copy(ret, list)
	m.listPool.Put(list)
	return ret
}

// Count 当前 ID 数量
func (m *Manager) Count() (count int) {
	m.data.Range(func(_, _ interface{}) bool {
		count++
		return true
	})
	return
}

// Clean 手动清理过期 ID；返回实际删除条数
func (m *Manager) Clean(nowTime time.Time, ttl time.Duration) int {
	now := nowTime.UnixNano()
	var del int
	m.data.Range(func(k, v interface{}) bool {
		node := v.(*idNode)
		if node.expired(now, ttl) {
			m.data.Delete(k)
			del++
		}
		return true
	})
	return del
}
